﻿using System;
using System.Runtime.InteropServices;
using Microsoft.Win32;

namespace Tsanie.ShellExt
{

    #region - Custom Attribute -

    [AttributeUsage(AttributeTargets.Class, AllowMultiple = false, Inherited = false)]
    public sealed class ExtensionFileTypesAttribute : Attribute
    {
        private bool m_findAssoc = false;
        private string[] m_FileTypes;

        public ExtensionFileTypesAttribute(string fileTypes) { m_FileTypes = fileTypes.Split(','); }
        public string FileTypes { get { return String.Join(",", m_FileTypes); } }
        public string[] FileTypesArray { get { return m_FileTypes; } }
        public bool FindAssoc
        {
            get { return m_findAssoc; }
            set { m_findAssoc = value; }
        }
    }

    #endregion

    public class Common
    {

        #region - APIs -

        [DllImport("kernel32.dll", CharSet = CharSet.Auto)]
        public static extern void OutputDebugString(string s);
        [DllImport("Shell32.dll")]
        public static extern void SHChangeNotify(int wEventId, uint uFlags, IntPtr dwItem1, IntPtr dwItem2);
        public const int SHCNE_ASSOCCHANGED = 0x08000000;
        [DllImport("uxtheme.dll", CharSet = System.Runtime.InteropServices.CharSet.Unicode)]
        public extern static int SetWindowTheme(IntPtr hWnd, string pszSubAppName, string pszSubIdList);

        #endregion

        #region - Debug Method -

        public static void OutputDebugString(string format, object arg1)
        {
            OutputDebugString(format, new object[] { arg1 });
        }
        public static void OutputDebugString(string format, object arg1, object arg2)
        {
            OutputDebugString(format, new object[] { arg1, arg2 });
        }
        public static void OutputDebugString(string format, object arg1, object arg2, object arg3)
        {
            OutputDebugString(format, new object[] { arg1, arg2, arg3 });
        }
        public static void OutputDebugString(string format, params object[] args)
        {
            OutputDebugString(string.Format(format, args));
        }

        public static void OutputError(string title, Exception e)
        {
            Common.OutputDebugString("[" + title + "] =======================");
            Common.OutputDebugString(e.GetType().FullName + ", " + e.Message);
            Common.OutputDebugString(e.StackTrace);
            Common.OutputDebugString("-----------------------------------------");
        }

        #endregion

        #region - Register/Unregister Method -

        static string CreateExtensionKey(string assoc, string extensionKey, string guid)
        {
            string result = assoc + "\\shellex\\" + extensionKey;
            try
            {
                using (RegistryKey key = Registry.ClassesRoot.CreateSubKey(result))
                {
                    key.SetValue(null, guid);
                    key.Close();
                }
                return result;
            }
            catch (Exception e)
            {
                Common.OutputDebugString("CreateExtensionKey():: Exception, '{0}', {1} - {2}", result, e.GetType().FullName, e.Message);
            }
            return null;
        }
        static string DeleteExtensionKey(string key)
        {
            try
            {
                Registry.ClassesRoot.DeleteSubKeyTree(key);
                return key;
            }
            catch (Exception e)
            {
                Common.OutputDebugString("DeleteExtensionKey(string):: Exception, '{0}', {1} - {2}", key, e.GetType().FullName, e.Message);
            }
            return null;
        }
        static string GetClassGuid(Type t)
        {
            GuidAttribute GuidAtt;
            try
            {
                GuidAtt = (GuidAttribute)(t.GetCustomAttributes(typeof(GuidAttribute), true)[0]);
                return "{" + GuidAtt.Value + "}";
            }
            catch (Exception e)
            {
                if (t == null)
                    Common.OutputDebugString("GetClassGuid(Type):: Argument 't' is null. - {1}", e.Message);
                else
                    Common.OutputDebugString("GetClassGuid(Type):: '{0}' don't have GuidAttribute. - {1}", t.FullName, e.Message);
            }
            return null;
        }

        internal static void RegisterExtension(Type t, string extensionKey)
        {
            string guid = GetClassGuid(t);
            if (string.IsNullOrEmpty(guid))
            {
                Common.OutputDebugString("RegisterExtension(Type, string):: [register] failed. Guid of '{0}' is null or empty.", t.FullName);
                return;
            }
            using (RegistryKey key = Registry.LocalMachine.CreateSubKey(@"Software\Microsoft\Windows\CurrentVersion\Shell Extensions\Approved"))
            {
                key.SetValue(guid, t.FullName);
                key.Close();
                //Common.OutputDebugString(@"RegisterExtension(Type, string):: Register extension approved - 'HKLM\Software\Microsoft\Windows\CurrentVersion\Shell Extensions\Approved, {0} = {1}'", guid, t.FullName);
            }
            try
            {
                object[] attrs = t.GetCustomAttributes(typeof(ExtensionFileTypesAttribute), true);
                if (attrs != null && attrs.Length > 0)
                {
                    ExtensionFileTypesAttribute fileTypes = (ExtensionFileTypesAttribute)attrs[0];
                    foreach (string type in fileTypes.FileTypesArray)
                    {
                        string result = null;
                        if (fileTypes.FindAssoc)
                        {
                            // 按照优先级依次查找关联
                            using (RegistryKey fileExt = Registry.CurrentUser.OpenSubKey(@"Software\Microsoft\Windows\CurrentVersion\Explorer\FileExts\" + type + @"\UserChoice"))
                            {
                                if (fileExt != null)
                                {
                                    string assocName = fileExt.GetValue("Progid") as string;
                                    fileExt.Close();
                                    if (assocName != null)
                                    {
                                        result = CreateExtensionKey(assocName, extensionKey, guid);
                                    }
                                }
                                if (result == null)
                                {
                                    // 查找HKCL\(ext)
                                    using (RegistryKey ext = Registry.ClassesRoot.OpenSubKey(type))
                                    {
                                        if (ext != null)
                                        {
                                            string assocName = ext.GetValue(null) as string;
                                            ext.Close();
                                            if (assocName != null)
                                            {
                                                result = CreateExtensionKey(assocName, extensionKey, guid);
                                            }
                                        }
                                    }
                                }
                            }
                        }
                        if (result == null)
                        {
                            // 如果都没有，可能是新建类型
                            result = CreateExtensionKey(type, extensionKey, guid);
                        }
                        //Common.OutputDebugString("RegisterExtension(Type, string):: [register] successfully. '{0}'", result);
                    }

                    // Tell Explorer to refresh
                    //Common.SHChangeNotify(Common.SHCNE_ASSOCCHANGED, 0, IntPtr.Zero, IntPtr.Zero);
                }
            }
            catch (Exception e)
            {
                Common.OutputDebugString("RegisterExtension(Type, string):: [register] exception: {0} - {1}", e.GetType().FullName, e.Message);
            }
        }
        internal static void UnregisterExtension(Type t, string extensionKey)
        {
            string guid = GetClassGuid(t);
            if (string.IsNullOrEmpty(guid))
            {
                Common.OutputDebugString("UnregisterExtension(Type, string):: [unregister] failed. Guid of '{0}' is null or empty.", t.FullName);
                return;
            }
            using (RegistryKey key = Registry.LocalMachine.CreateSubKey(@"Software\Microsoft\Windows\CurrentVersion\Shell Extensions\Approved"))
            {
                key.DeleteValue(guid);
                key.Close();
                //Common.OutputDebugString(@"UnregisterExtension(Type, string):: Unregister extension approved - 'HKLM\Software\Microsoft\Windows\CurrentVersion\Shell Extensions\Approved, {0} = {1}'", guid, t.FullName);
            }
            try
            {
                object[] attrs = t.GetCustomAttributes(typeof(ExtensionFileTypesAttribute), true);
                if (attrs != null && attrs.Length > 0)
                {
                    ExtensionFileTypesAttribute fileTypes = (ExtensionFileTypesAttribute)attrs[0];
                    foreach (string type in fileTypes.FileTypesArray)
                    {
                        string result = null;
                        if (fileTypes.FindAssoc)
                        {
                            // 按照优先级依次查找关联
                            using (RegistryKey fileExt = Registry.CurrentUser.OpenSubKey(@"Software\Microsoft\Windows\CurrentVersion\Explorer\FileExts\" + type + @"\UserChoice"))
                            {
                                if (fileExt != null)
                                {
                                    string assocName = fileExt.GetValue("Progid") as string;
                                    fileExt.Close();
                                    if (assocName != null)
                                    {
                                        result = DeleteExtensionKey(assocName + @"\shellex\" + extensionKey);
                                    }
                                }
                                if (result == null)
                                {
                                    // 查找HKCL\(ext)
                                    using (RegistryKey ext = Registry.ClassesRoot.OpenSubKey(type))
                                    {
                                        if (ext != null)
                                        {
                                            string assocName = ext.GetValue(null) as string;
                                            ext.Close();
                                            if (assocName != null)
                                            {
                                                result = DeleteExtensionKey(assocName + @"\shellex\" + extensionKey);
                                            }
                                        }
                                    }
                                }
                            }
                        }
                        if (result == null)
                        {
                            // 如果都没有，可能是新建类型
                            result = DeleteExtensionKey(type + @"\shellex\" + extensionKey);
                        }
                        //Common.OutputDebugString("UnregisterExtension(Type, string):: [unregister] successfully. '{0}'", result);
                    }

                    // Tell Explorer to refresh
                    Common.SHChangeNotify(Common.SHCNE_ASSOCCHANGED, 0, IntPtr.Zero, IntPtr.Zero);
                }
            }
            catch (Exception e)
            {
                Common.OutputDebugString("UnregisterExtension(Type, string):: [unregister] exception - {0}", e.Message);
            }
        }

        #endregion

    }
}
